from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)


# 计算雅克比矩阵
"""
W_ir = params["params"]["GRUCell_0"]["ir"]["kernel"]
W_iz = params["params"]["GRUCell_0"]["iz"]["kernel"]
W_in = params["params"]["GRUCell_0"]["in"]["kernel"]
W_hr = params["params"]["GRUCell_0"]["hr"]["kernel"]
W_hz = params["params"]["GRUCell_0"]["hz"]["kernel"]
W_hn = params["params"]["GRUCell_0"]["hn"]["kernel"]
b_in = params["params"]["GRUCell_0"]["in"]["bias"]
b_hr = params["params"]["GRUCell_0"]["hr"]["bias"]
b_hz = params["params"]["GRUCell_0"]["hz"]["bias"]
b_hn = params["params"]["GRUCell_0"]["hn"]["bias"]
"""
def gru_step(x, h, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hn):
    # Directly use leaves by index
    # Reset gate
    r = jax.nn.sigmoid(jnp.dot(x, W_ir) + jnp.dot(h, W_hr))
    
    # Update gate
    z = jax.nn.sigmoid(jnp.dot(x, W_iz) + jnp.dot(h, W_hz))
    
    # Candidate hidden state
    n = jnp.tanh(jnp.dot(x, W_in) + b_in + r * (jnp.dot(h, W_hn) + b_hn))
    
    # New hidden state
    h_new = (1 - z) * n + z * h
    
    return h_new

@jax.jit
def rnn_run_vector(state, obs, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hn):
    state1 = gru_step(obs, state, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hn)
    return state1 - state

jacobian_fun = jax.jacrev(rnn_run_vector)
jacobian_fun_batch = jax.vmap(jacobian_fun, in_axes=(0, None, None, None, None, None, None, None, None, None))

def compute_eig(sample_point, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hr, b_hz, b_hn):
    obs = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    jacobian_of_x = jacobian_fun(jnp.array(sample_point), obs, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hr, b_hz, b_hn)
    eigenvalues = np.linalg.eigvals(jacobian_of_x)
    if any(np.real(eigenvalues) > 0):
        return False
    else:
        return True
    
@jax.jit
def check_diverge(eigenvalues):
    return jnp.any(jnp.real(eigenvalues) > 0)

check_diverge_batch = jax.vmap(check_diverge)

def compute_eig_batch(sample_points, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hn):
    obs = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    jacobian_of_x = jacobian_fun_batch(sample_points, obs, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hn)
    eigenvalues = np.linalg.eigvals(jacobian_of_x)
    
    # 批量统计是否有发散的情况
    diverge_batch = check_diverge_batch(eigenvalues)
    return diverge_batch


def main():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    """
    r = \sigma(W_{ir} x + W_{hr} h + b_{hr})
    z = \sigma(W_{iz} x + W_{hz} h + b_{hz})
    n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn}))
    h' = (1 - z) * n + z * h
    """
    # print(params["params"]["GRUCell_0"]["hn"]["kernel"].shape)
    W_ir = params["params"]["GRUCell_0"]["ir"]["kernel"]
    W_iz = params["params"]["GRUCell_0"]["iz"]["kernel"]
    W_in = params["params"]["GRUCell_0"]["in"]["kernel"]
    W_hr = params["params"]["GRUCell_0"]["hr"]["kernel"]
    W_hz = params["params"]["GRUCell_0"]["hz"]["kernel"]
    W_hn = params["params"]["GRUCell_0"]["hn"]["kernel"]
    b_in = params["params"]["GRUCell_0"]["in"]["bias"]
    b_hn = params["params"]["GRUCell_0"]["hn"]["bias"]
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 1000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    for t in range(10):
        progress_bar(t, rpl_config.life_duration)
        """ model forward 
        """
        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)

    rnn_state_1 = jnp.copy(rnn_state)

    rnn_state_trajectory = []

    for t in range(rpl_config.life_duration):
        progress_bar(t, rpl_config.life_duration)
        """ model forward 
        """
        # # test gru_step
        # rnn_state_1 = gru_step(obs_zero, rnn_state, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hn)

        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        rnn_state_trajectory.append(np.array(rnn_state).copy())

        # print("diff between rnn_state_1 and rnn_state: ", np.linalg.norm(rnn_state_1 - rnn_state))
            
    # # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    # rnn_state_trajectory1_np = np.array(rnn_state_trajectory)
    # # 交换 rnn_state_trajectory_np 的第一维和第二维
    # # rnn_state_trajectory1_np = np.swapaxes(rnn_state_trajectory1_np, 0, 1)
    # print("shape of rnn_state_trajectory1_np: ", rnn_state_trajectory1_np.shape)

    # diverge_batch = compute_eig_batch(rnn_state_trajectory1_np[-1], W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hn)
    
    # # 统计 diverge_batch 有多少个 True和False
    # n_diverge = jnp.sum(diverge_batch)

    # print("n_diverge ratio: ", n_diverge / diverge_batch.shape[0])

    # def calculate_divergence(F, x, epsilon=1e-5):
    #     """
    #     计算离散动力系统在点x处的散度
        
    #     参数:
    #     F : 函数, 接受一个numpy数组作为输入,返回下一个状态
    #     x : numpy数组, 计算散度的位置
    #     epsilon : float, 用于数值微分的小增量
        
    #     返回:
    #     float: 在点x处估计的散度
    #     """
    #     n = len(x)
    #     divergence = 0.0
        
    #     for i in range(n):
    #         # 创建单位向量
    #         e = np.zeros(n)
    #         e[i] = 1.0
            
    #         # 计算F在x + epsilon * e 和 x - epsilon * e 处的值
    #         f_plus = F(x + epsilon * e)
    #         f_minus = F(x - epsilon * e)
            
    #         # 使用中心差分法估算偏导数
    #         partial_derivative = (f_plus[i] - f_minus[i]) / (2 * epsilon)
            
    #         divergence += partial_derivative
        
    #     # print(divergence)
    #     return divergence

    # def F(state):
    #     return gru_step(obs_zero, state, W_ir, W_iz, W_in, W_hr, W_hz, W_hn, b_in, b_hr, b_hz, b_hn)

    # for state in rnn_state_trajectory1_np[-1]:
    #     print("shape of state: ", state.shape)
    #     diver = calculate_divergence(F, state)
    #     print("divergence: ", np.sum(diver))

    # # 生成和 rnn_state 形状相同的 noise
    # k2 = jax.random.PRNGKey(np.random.randint(0, 1000000))
    # epsilon = 0.01
    # noise = jax.random.uniform(k2, shape = rnn_state.shape, minval = -epsilon, maxval = epsilon)
    # rnn_state = rnn_state_1 + noise

    # rnn_state_trajectory = []

    # for t in range(rpl_config.life_duration):
    #     progress_bar(t, rpl_config.life_duration)
    #     """ model forward 
    #     """
    #     rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
    #     rnn_state_trajectory.append(np.array(rnn_state).copy())

    # rnn_state_trajectory2_np = np.array(rnn_state_trajectory)
    # # rnn_state_trajectory2_np = np.swapaxes(rnn_state_trajectory2_np, 0, 1)
    # print("shape of rnn_state_trajectory2_np: ", rnn_state_trajectory2_np.shape)

    # # 计算两批轨迹的距离
    # distance = np.linalg.norm(rnn_state_trajectory1_np - rnn_state_trajectory2_np, axis = 2)
    # distance = distance.T
    # print("shape of distance: ", distance.shape)

    # def lyapunov_exp(dist):
    #     # 检查 dist 数组中是否存在零值
    #     if 0 in dist:
    #         # 复制 dist 数组，将零值替换为一个非零的小值（例如 1e-10）
    #         dist = np.where(dist == 0, 1e-10, dist)
    #     lnr = np.log(dist)
    #     slope, intercept = np.polyfit(list(range(len(lnr))), lnr, deg=1)
    #     return slope
    
    # lyapunov_exp_list = []
    # for i in range(1, distance.shape[0]):
    #     progress_bar(i, distance.shape[0])
    #     lyapunov_exp_list.append(lyapunov_exp(distance[i]))
    # lyapunov_exp_list = np.array(lyapunov_exp_list)
    # print("shape of lyapunov_exp_list: ", lyapunov_exp_list.shape)

    # # 用直方图统计 lyapunov_exp_list 的分布
    # hist, bin_edges = np.histogram(lyapunov_exp_list, bins = 100)
    # max_idx = np.argmax(hist)
    # # 绘制直方图
    # fig = plt.figure()
    # plt.hist(lyapunov_exp_list, bins = 100)
    # plt.axvline(x = bin_edges[max_idx], color = 'r')
    # plt.text(bin_edges[max_idx], 25, str(bin_edges[max_idx]), fontsize = 20)
    # plt.show()


if __name__ == "__main__":
    main()